#include "dag.h"
#include "dag_operators.h"
#include "dag_types.h"

#include "kdtree.h"

#include <string>
#include <memory>
#include <set>

using namespace dag;
using namespace std;

KdTreeOp::KdTreeOp()
:	m_InputSpec(FrameSpec()),
	m_OutputSpec(FrameSpec()),
	m_Tree(NULL),
	m_UseFloat32(false)
{
}

KdTreeOp::KdTreeOp(const FrameDepot& inputs,
				   const FrameDepot& outputs)
:	m_InputSpec(inputs.spec()),
	m_OutputSpec(outputs.spec()),
	m_Tree(NULL),
	m_UseFloat32(false)
{
	int k = 0;
	if (m_InputSpec.fSize > 0) {
		k = m_InputSpec.fSize;
		m_UseFloat32 = true;
		if (m_InputSpec.dSize > 0) {
			throw InvalidArgument();
		}
	}
	else if (m_InputSpec.dSize == 0) {
		throw InvalidArgument();
	}
	else {
		k = m_InputSpec.dSize;
		m_UseFloat32 = false;
	}
	
	if (inputs.size() != outputs.size()) {
		throw InvalidArgument();
	}
	
	m_Tree = reinterpret_cast<void*> (kd_create(k));
	if (m_Tree == NULL) {
		throw NullPtr();
	}
	
	for (size_t i = 0; i < inputs.size(); i++) {
		auto_ptr<Frame> newInFrame(new Frame(m_InputSpec));
		auto_ptr<Frame> newOutFrame(new Frame(m_OutputSpec));
		
		inputs.fill(i, *newInFrame);
		outputs.fill(i, *newOutFrame);
		addPair(newInFrame, newOutFrame);
	}
}

KdTreeOp::KdTreeOp(istream& bStream) 
:	m_Tree(NULL),
	m_UseFloat32(false)
{
	int kval = -1;
	int fbits = -1;
	unsigned numFrames = -1;
	
	bStream.read((char*)&fbits, dag::dag_int32_size);
	bStream.read((char*)&kval, dag::dag_int32_size);
	bStream.read((char*)&numFrames, dag::dag_uint32_size);
	
	if (fbits != 32 && fbits != 64) {
		throw InvalidValue();
	}
	if (kval <= 0) {
		throw InvalidValue();
	}
	if (numFrames == 0) {
		throw InvalidValue();
	}
	
	/* set data type of position coordinates */
	if (fbits == 32) {
		m_UseFloat32 = true;
	}
	else {
		m_UseFloat32 = false;
	}
	
	/* create tree */
	m_Tree = reinterpret_cast<void*> (kd_create(kval));
	if (m_Tree == NULL) {
		throw NullPtr();
	}
	
	/* load input/output frames */
	for (unsigned i = 0; i < numFrames; i++) {
		auto_ptr<Frame> input = Frame::inflate(bStream);
		auto_ptr<Frame> output = Frame::inflate(bStream);
		addPair(input, output);
	}
}

KdTreeOp::~KdTreeOp() throw() {
	set<Frame*>::iterator setIter;
	for (setIter = m_InputFrames.begin(); 
		 setIter != m_InputFrames.end(); 
		 setIter++) {
		delete *setIter;
	}
	
	for (setIter = m_OutputFrames.begin(); 
		 setIter != m_OutputFrames.end(); 
		 setIter++) {
		delete *setIter;
	}
	
	kd_free(reinterpret_cast<kdtree*>(m_Tree));
}

auto_ptr<FrameOperator> KdTreeOp::inflate(istream& bStream) {
	return auto_ptr<FrameOperator>(new KdTreeOp(bStream));
}

void KdTreeOp::operate(const Frame& input, Frame& output) {
	kdres* searchResult = NULL;
	
	if (m_UseFloat32) {
		searchResult = kd_nearestf(reinterpret_cast<kdtree*>(m_Tree),
								   input.floats);
	}
	else {
		searchResult = kd_nearest(reinterpret_cast<kdtree*>(m_Tree),
								  input.doubles);
	}
	
	if (kd_res_size(searchResult) == 1) {
		output = Frame(m_OutputSpec);
	}
	
	const Frame* frame = NULL;
	frame = reinterpret_cast<Frame*>(kd_res_item_data(searchResult));
	if (frame == NULL) {
		output = Frame(m_OutputSpec);
	}
	else {
		output = *frame;
	}
}


void KdTreeOp::addPair(auto_ptr<Frame> input, auto_ptr<Frame> output) {
	pair<set<Frame*>::iterator, bool> m_InsertResult1;
	pair<set<Frame*>::iterator, bool> m_InsertResult2;

	m_InsertResult1 = m_InputFrames.insert(input.release());
	m_InsertResult2 = m_OutputFrames.insert(output.release());
		
	if (m_UseFloat32) {
		kd_insertf(reinterpret_cast<kdtree*>(m_Tree),
				   (*m_InsertResult1.first)->floats,
				   *m_InsertResult2.first);
	}
	else {
		kd_insert(reinterpret_cast<kdtree*>(m_Tree),
				  (*m_InsertResult1.first)->doubles,
				  *m_InsertResult2.first);
	}
}

const string& KdTreeOp::_frameOperatorID() const {
	return OpId::kd_tree_op;
}

void KdTreeOp::_serialize(ostream& binaryStream) const {
	/* serialize kdtree - saves the k value, number of float bits and the size
	 * of the input outputs.  Cannot actually save entire tree structure, so we
	 * just save the values and rebuild the tree upon inflation */
	int kval = -1;
	int fbits = -1;
	if (m_UseFloat32) {
		kval = m_InputSpec.fSize;
		fbits = 32;
	}
	else {
		kval = m_InputSpec.dSize;
		fbits = 64;
	}

	unsigned numFrames = m_InputFrames.size();
	
	binaryStream.write((const char*)&fbits, dag::dag_int32_size);
	binaryStream.write((const char*)&kval, dag::dag_int32_size);
	binaryStream.write((const char*)&numFrames, dag::dag_uint32_size);
	
	set<Frame*>::const_iterator fIter1 = m_InputFrames.begin();
	set<Frame*>::const_iterator fIter2 = m_OutputFrames.begin();
	while (fIter1 != m_InputFrames.end() && fIter2 != m_OutputFrames.end()) {		
		(*fIter1)->serialize(binaryStream);
		(*fIter2)->serialize(binaryStream);
		fIter1++;
		fIter2++;
	}
}
